import hmac
import json
import requests
import urllib.parse as urlparse
import hashlib
import sys
import time
import base64

from coin.exchange.base.kr_rest.auth_base import AuthKey
from coin2.api_client.common import (SyncQueryer)
from enum import Enum

class Type(Enum):
  BUY = "buy"
  SELL = "sell"

class OrderType(Enum):
  MARKET = "market"
  LIMIT = "limit"
  STOP_LOSS = "stop-loss"
  TAKE_PROFIT = "take-profit"
  STOP_LOSS_LIMIT = "stop-loss-limit"
  TAKE_PROFIT_LIMIT = "take-profit-limit"
  SETTLE_POSITION = "settle-position"

class TimeInForce(Enum):
  GTC = "GTC"
  IOC = "IOC"
  GTD = "GTD"

class KrakenAuth(requests.auth.AuthBase):
  def __init__(self, key):
    self._key = AuthKey.from_file(key)

  def _gen_signature(self, urlpath: str, nonce: str, body: str):
    assert urlpath, "urlpath must be not None and not empty!"
    assert nonce, "nonce must be not None and not empty!"
    assert body, "body must be not None and not empty!"

    encoded = (nonce + body).encode()
    message = urlpath.encode() + hashlib.sha256(encoded).digest()

    mac = hmac.new(base64.b64decode(self._key.secret_key), message, hashlib.sha512)
    sigdigest = base64.b64encode(mac.digest())
    return sigdigest.decode()

  def __call__(self, req: requests.PreparedRequest):
    nonce = str(int(time.time() * 1000))
    path = urlparse.urlparse(req.url).path
    
    body = dict(urlparse.parse_qsl(req.body))
    body['nonce'] = nonce
    req.prepare_body(body, None)

    signature = self._gen_signature(path, nonce, req.body)

    req.headers.update({'API-Key': self._key.access_key})
    req.headers.update({'API-Sign': signature})
    req.prepare_headers(req.headers)

    # print(req.headers)
    # print(req.body)
    # print(req.url)

    return req

class KrakenPrivateClient():
  def __init__(self, key_path: str, ip=None):
    self._auth = KrakenAuth(key_path)
    self._url = "https://api.kraken.com"

    self._query = SyncQueryer(api_host=self._url, auth=self._auth, ip=ip)

  def get_account_balance(self):
    method = "POST"
    data = {
    }
    path = "/0/private/Balance"

    response = self._query.query(method=method, path=path, data=data)
    return response.json()

  def get_open_order(self, trades: bool = None):
    method = "POST"
    data = {
      'trades': None if trades == None else "true" if trades else "false"
    }
    path = "/0/private/OpenOrders"

    response = self._query.query(method=method, path=path, data=data)
    return response.json()

  def cancel_order(self, txid: str):
    assert txid, "txid must be not None and not empty!"

    method = "POST"
    data = {
      'txid': txid,
    }
    path = "/0/private/CancelOrder"

    response = self._query.query(method=method, path=path, data=data)
    return response.json()

  def submit_order(
    self,
    order_type: OrderType,
    type: Type,
    volume: str,
    pair: str,
    price: str = None,
    time_in_force : TimeInForce = None
  ):
    assert order_type, "order_type must be not None!"
    assert type, "type must be not None!"
    assert volume, "volume must be not None and not empty!"
    assert pair, "pair must be not None and not empty!"

    method = "POST"
    data = {
      'ordertype': order_type.value,
      'type': type.value,
      'volume': volume,
      'pair': pair,
      'price': price,
      'timeinforce': time_in_force and time_in_force.value,
    }
    path = "/0/private/AddOrder"

    response = self._query.query(method=method, path=path, data=data)
    return response.json()

# api doc url
# https://docs.kraken.com/rest/

if __name__ == "__main__":
  key_path = sys.argv[1]
  
  kraken = KrakenPrivateClient(key_path)
  info = kraken.get_open_order()
  # info = kraken.get_account_balance()
  # info = kraken.cancel_order("1234")
  # info = kraken.submit_order(OrderType.LIMIT, Type.SELL, "0.00000001", "BTC/USD", "100000000")
  print(json.dumps(info, indent=2))
